import numpy as np
import matplotlib.pyplot as plt


def DLS_via_normal_equation(x, y):
    if np.size(x) != np.size(y):
        print('ERROR: the size of two vectors does not match!')
        return
    G = np.zeros((3, 3))
    c = np.zeros(3)
    for i in range(3):
        for j in range(3):
            k = i+j
            G[i][j] = sum(np.power(x, k))
    for i in range(3):
        c[i] = sum(y*np.power(x, i))
    a = np.linalg.solve(G, c)
    print('cond(G): ', np.linalg.cond(G))
    return a


def DLS_via_QR_factorization(x, y):
    if np.size(x) != np.size(y):
        print('ERROR: the sizes of two vectors do not match!')
        return
    n = np.size(x)
    A = np.zeros((n, 3))
    b = np.zeros(n)
    for i in range(n):
        for j in range(3):
            A[i][j] = pow(x[i], j)
        b[i] = y[i]
    Q, R = np.linalg.qr(A)
    c = np.dot(np.transpose(Q), b)
    c = c[0:3]
    a = np.linalg.solve(R[0:3], c)
    print('cond(R): ', np.linalg.cond(R))
    return a


def output(x, y):
    coe1 = DLS_via_normal_equation(x, y)
    coe2 = DLS_via_QR_factorization(x, y)
    print('DLS via normal equations coefficients: ', coe1)
    print('DLS via QR factorization coefficients: ', coe2)
    datanum = np.size(x)
    N = 1000
    points = np.zeros(N+1)
    value1 = np.zeros(N+1)
    value2 = np.zeros(N+1)
    for i in range(N+1):
        points[i] = x[0]+(x[datanum-1]-x[0])/N*i
        value1[i] = coe1[0] + coe1[1]*points[i] + coe1[2]*points[i]*points[i]
        value2[i] = coe2[0] + coe2[1] * points[i] + coe2[2] * points[i] * points[i]
    plt.xlabel('x value')
    plt.ylabel('y value')
    plt.plot(points, value1, color='y')
    plt.plot(points, value2, color='g')
    plt.legend(('DLS via normal equations', 'DLS via QR factorization'), loc='upper left')
    plt.scatter(x, y, color='r')

    global pic_num
    pic_name = 'DLS_for_data_' + str(pic_num)
    plt.savefig('./pic/' + pic_name)
    plt.show()

    pic_num += 1


pic_num = 1
x1 = [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]
y1 = [2.9, 2.7, 4.8, 5.3, 7.1, 7.6, 7.7]
print('data 1: ')
output(x1, y1)

x2 = [3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5]
y2 = [7.6, 9.4, 9.0, 9.6, 10.0, 10.2, 9.7]
print('data 2: ')
output(x2, y2)

x3 = [7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0]
y3 = [8.3, 8.4, 9.0, 8.3, 6.6, 6.7, 4.1]
print('data 3: ')
output(x3, y3)
